/*
 * extract_face_main.cpp
 *
 *  Created on: Mar 30, 2011
 *      Author: tqlong
 */

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <boost/filesystem/operations.hpp>
#include <cassert>
#include <iostream>
#include <fstream>
#include <algorithm>

#include <opencv2/opencv.hpp>
#include <armadillo>

#include "common/armavec.h"
#include "reginfomax/kmean.h"
//#include <opencv2/highgui.h>

namespace po = boost::program_options;
namespace fs = boost::filesystem;

po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);

typedef std::vector<cv::Rect> rvec_type;

rvec_type extract_face(const cv::Mat& img)
{
  std::string fname = vm["cascade"].as<std::string>();
  rvec_type r;
  cv::CascadeClassifier cc(fname);
  if (cc.empty()) {
    std::cerr << "Check cascade ..." << fname << "\n";
    return r;
  }
  cc.detectMultiScale(img, r, 1.1, 3, 0, cv::Size(40,40));
//  for (unsigned int i = 0; i < r.size(); i++) {
//    r[i].x -= r[i].width/2;
//    r[i].y -=
//  }
  return r;
}

void draw_rect(cv::Mat& img, const rvec_type& r)
{
  for (unsigned int i = 0; i < r.size(); i++)
    cv::rectangle(img, r[i], cv::Scalar(0, 0, 255), 2);
}

void convertMat2Arma(const cv::Mat& face, arma::mat &face_vec)
{
  cv::Mat gray;
  cv::cvtColor(face, gray, CV_RGB2GRAY);
//  cv::convertScaleAbs(gray, dgray, 1. / 255);
//  std::cerr << "dgray.type = " << dgray.type() << "\n";
  face_vec = arma::zeros(gray.rows, gray.cols);
  for (int i = 0; i < gray.rows; i++)
    for (int j = 0; j < gray.cols; j++)
      face_vec.at(i, j) = (double) gray.at<unsigned char>(j, i) / 255.0;
}

class ImageDataset : public Dataset<ArmaVec>
{
public:
  typedef Dataset<ArmaVec>           __Base;
  typedef __Base::vec_type           vec_type;
public:
  std::vector<std::string> fname_;

  void load(const std::string& dir) {
    fs::path face_dirpath(dir);
    fs::directory_iterator itr(face_dirpath), end_itr;
    for (; itr != end_itr; itr++) {
      fs::path face_path(face_dirpath / itr->leaf());
      cv::Mat face = cv::imread(face_path.relative_path().string());
      if (face.empty()) continue;

      std::cerr << face_path.leaf() << " : " << face.cols << " x " << face.rows << "\n";
      vec_type face_vec;
      convertMat2Arma(face, face_vec);
      push_back(face_vec);
      fname_.push_back(face_path.leaf());
    }
  }

  const std::string& name(unsigned int i) const {
    assert(i < size());
    return fname_[i];
  }

  static void convertMat2Arma(const cv::Mat& face, vec_type &face_vec)
  {
    cv::Mat gray, scale, ehist;
    cv::cvtColor(face, gray, CV_RGB2GRAY);
    cv::resize(gray, scale, cv::Size(32,32));
    cv::equalizeHist(scale, ehist);
  //  cv::convertScaleAbs(gray, dgray, 1. / 255);
  //  std::cerr << "dgray.type = " << dgray.type() << "\n";
    face_vec = arma::zeros(scale.rows * scale.cols);
    int k = 0;
    for (int j = 0; j < scale.cols; j++)
      for (int i = 0; i < scale.rows; i++, k++)
        face_vec[k] = (double) ehist.at<unsigned char>(j, i) / 255.0;
  }
};

class EigenDataset : public Dataset<ArmaVec>
{
public:
  typedef Dataset<ArmaVec>           __Base;
  typedef __Base::vec_type           vec_type;
public:
  std::vector<std::string> fname_;
  arma::mat U_;

  void load(const ImageDataset& data) {
    arma::mat A(data.dim(), data.n()), U, V;
    std::cerr << "A.rows = " << A.n_rows << " A.cols = " << A.n_cols << "\n";
    arma::vec s;
    for (int i = 0; i < data.n(); i++)
      A.col(i) = data[i];
    arma::svd(U, s, V, A);
    U_ = U.cols(0, 9);
    std::cerr << "U.rows = " << U.n_rows << " U.cols = " << U.n_cols << "\n";
    std::cerr << "U_.rows = " << U_.n_rows << " U_.cols = " << U_.n_cols << "\n";
    for (int i = 0; i < data.n(); i++) {
      arma::vec feature = arma::trans(U_)*data[i];
      push_back(feature);
      fname_.push_back(data.name(i));
    }
  }

  const std::string& name(unsigned int i) {
    assert(i < size());
    return fname_[i];
  }

};


int main(int argc, char** argv) {
  process_options(argc, argv);
  if (vm.count("help")) {
    std::cout << desc << "\n";
    return 0;
  }

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  if (vm.count("file")) {
    cv::Mat img;
    img = cv::imread(vm["file"].as<std::string>());
    std::cerr << vm["file"].as<std::string>() << " : " << img.cols << " x " << img.rows << "\n";
    rvec_type rect = extract_face(img);
    draw_rect(img, rect);
    cv::imshow("extract", img);
    cv::waitKey(0);
    return 0;
  }

  if (vm.count("dir")) {
    std::string dname = vm["dir"].as<std::string>();
    fs::path dirpath(dname), face_dirpath(dirpath / "faces");
    fs::create_directory(face_dirpath);
    fs::directory_iterator itr(dirpath), end_itr;
    int count = 0;
    for (; itr != end_itr; itr++) {
      std::string fname(itr->leaf());
      fs::path fpath(dirpath / fname);
      if (fname.substr(fname.length()-4) != ".jpg") continue;
      count++;
      cv::Mat img = cv::imread(fpath.relative_path().string());
//      std::cerr << "path = " << fpath.relative_path().string() << "\n";
      std::cerr << fpath.leaf() << " : " << img.cols << " x " << img.rows << "\n";
      rvec_type rect = extract_face(img);
      //draw_rect(img, rect);
      //cv::imshow("extract", img);
      for (unsigned int i = 0; i < rect.size(); i++) {
        cv::Mat roi = img(rect[i]), size128;
        cv::resize(roi, size128, cv::Size(128,128));
        char c = '0'+i;
        std::stringstream ss;
        ss << fname << "_" << c << ".jpg";
        std::string facename = ss.str();
        fs::path facepath(face_dirpath / facename);
        std::cerr << "path = " << facepath.relative_path().string() << "\n";
        cv::imwrite(facepath.relative_path().string(), size128);
      }
    }
    std::cerr << "total " << count << " images\n";
  }

//  if (vm.count("face") && vm.count("arma")) {
//    fs::path face_dirpath(vm["face"].as<std::string>());
//    fs::path arma_dirpath(vm["arma"].as<std::string>());
//    fs::create_directory(arma_dirpath);
//    fs::directory_iterator itr(face_dirpath), end_itr;
//    for (; itr != end_itr; itr++) {
//      fs::path face_path(face_dirpath / itr->leaf()), arma_path(arma_dirpath / (itr->leaf() + ".arma"));
//      cv::Mat face = cv::imread(face_path.relative_path().string());
//      std::cerr << face_path.leaf() << " : " << face.cols << " x " << face.rows << "\n";
//      arma::mat face_vec;
//      convertMat2Arma(face, face_vec);
//      face_vec.save(arma_path.relative_path().string(), arma::arma_binary);
//    }
//  }

  if (vm.count("face")) {
    ImageDataset idata;
    idata.load(vm["face"].as<std::string>());
    idata.normalize();
    EigenDataset data;
    data.load(idata);

    int nc = vm["nc"].as<int>();
    int maxiter = vm["iter"].as<int>();
    std::vector<double> p;
    p.push_back(maxiter);

//    typedef KMean<ImageDataset::vec_type, ImageDataset> Functional;
    typedef KMean<EigenDataset::vec_type, EigenDataset> Functional;
    Functional F(data, nc);
    F.setParams(p);
    F.optimize();

    int n = data.size();
    std::vector<std::vector<std::string> > member(nc);
    for (int i = 0; i< n; i++) {
      std::cout << "sample " << i << ": " << data.name(i) << " class " << F.y(i) << "\n";
      member[F.y(i)].push_back(data.name(i));
    }

    std::ofstream f(vm["output"].as<std::string>().c_str());
    for (int k = 0; k < nc; k++) {
      f << "----- class " << k << "-----\n";
      std::sort(member[k].begin(), member[k].end());
      for (unsigned int i = 0; i < member[k].size(); i++) {
        f << member[k][i] << "\n";
      }
    }
    f.close();

//    for (int i = 0; i < nc; i++) {
//      std::cout << "center[" << i << "]=\n" << F.center(i) << "\n";
//    }

  }
  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  std::cerr << "Duration = " << duration << "\n";
  return 0;
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("file", po::value<std::string>(), "image file contains face")
    ("dir", po::value<std::string>(), "raw data")
    ("face", po::value<std::string>(), "face image dir")
    ("arma", po::value<std::string>(), "arma vector dir")
    ("output", po::value<std::string>()->default_value("member-output"), "membership output")
    ("cascade", po::value<std::string>()->default_value("/usr/local/share/opencv/haarcascades/haarcascade_frontalface_default.xml"), "cascade classifier xml")
//    ("gen", "generate random data")
//    ("match", "match two datasets")
//    ("method", po::value<std::string>()->default_value("auctree"), "matching method: auction, alternate, auctree")
//    ("dat1", po::value<std::string>(), "file contains the first dataset")
//    ("dat2", po::value<std::string>(), "file contains the second dataset")
//    ("seed", "using different random seed")
//    ("n", po::value<int>()->default_value(10), "number of points")
    ("nc", po::value<int>()->default_value(3), "number of centers")
//    ("d", po::value<int>()->default_value(2), "input dimension")
//    ("r", po::value<double>()->default_value(1.0), "cluster radius")
//    ("lambda", po::value<double>()->default_value(1.0), "regularized parameter")
    ("iter", po::value<int>()->default_value(10), "maximum number of iterations")
//    ("eta", po::value<double>()->default_value(0.5), "learning rate")
//    ("mesh", po::value<std::string>()->default_value("mesh"), "mesh files prefix")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}
